import os
from training.msg import ChunkMsg,GeneralMsg
from iop import BusinessOperation

import json
import os
import qianfan
import time
import intersystems_iris.dbapi._DBAPI as dbapi

class SaveToVector(BusinessOperation):
    
    def on_init(self):
        os.environ["QIANFAN_ACCESS_KEY"] = "a8bc3628280448d790fbb6573040eaa2"
        os.environ["QIANFAN_SECRET_KEY"] = "e9c76bf17ba44c379703666040c5b7c8"
        return

    def on_vector_request(self, request: ChunkMsg):
        config = {
        "hostname": "localhost",
        "port": 1972,
        "namespace": "IRISAPP",
        "username": "superuser",
        "password": "SYS",
        }
        connection = dbapi.connect(**config)
        cursor = connection.cursor()
        sql = "INSERT INTO MyIRIS.VectorLab (description, description_vector) " + "VALUES (?,TO_VECTOR(?,double))"
        # json_data = json.dumps(request.docs[0])
        all_splits = request.docs
        self.trace(str(all_splits[0]["page_content"]))
        self.trace("1")
        emb = qianfan.Embedding()  # 初始化嵌入模型对象
        #embeddings = []
        for chunk in all_splits:  # 遍历所有分割的文本块
        # 获取文本块的嵌入向量，使用默认模型Embedding-V1
            description = chunk["page_content"]
            self.trace(description)
            resp = emb.do(texts=[description], model="Embedding-V1")  # 请求嵌入向量           
            #embeddings.append(str(resp['data'][0]['embedding']))  # 将嵌入向量添加到列表中
            em = str(resp['data'][0]['embedding'])
            cursor.execute(sql, (description, em))
            time.sleep(1)  # 暂停1秒，避免请求过于频繁
        #self.trace(str(embeddings[0]))  # 打印第一个文本块的嵌入向量
        connection.close()           
        response = GeneralMsg()
        response.msg = 'OK'
        return response
            